import os
import time
import traceback
import cv2
import nibabel as nib
import argparse
import numpy as np
import SimpleITK as sitk
import logging

logging.basicConfig(level=logging.DEBUG)


# user_input -> nnunet_input
def write_2d_to_3d(from_path: str, to_path: str):
    logging.info("Start to convert ori data to nii")
    os.makedirs(to_path, exist_ok=True)

    file_list = os.listdir(from_path)
    t1 = time.perf_counter()
    for each_file in file_list:
        t2 = time.perf_counter()
        img = cv2.imread(os.path.join(from_path, each_file))
        img_transposed = np.transpose(img, (2, 0, 1))
        img_0 = img_transposed[0].reshape(1, img_transposed[0].shape[0], img_transposed[0].shape[1])
        img_1 = img_transposed[1].reshape(1, img_transposed[1].shape[0], img_transposed[1].shape[1])
        img_2 = img_transposed[2].reshape(1, img_transposed[2].shape[0], img_transposed[2].shape[1])

        name, _ = os.path.splitext(each_file)

        img_0_name = name + '_0000.nii.gz'
        img_1_name = name + '_0001.nii.gz'
        img_2_name = name + '_0002.nii.gz'

        img_0_nii = sitk.GetImageFromArray(img_0)
        img_1_nii = sitk.GetImageFromArray(img_1)
        img_2_nii = sitk.GetImageFromArray(img_2)

        sitk.WriteImage(img_0_nii, os.path.join(to_path, img_0_name))
        sitk.WriteImage(img_1_nii, os.path.join(to_path, img_1_name))
        sitk.WriteImage(img_2_nii, os.path.join(to_path, img_2_name))
        t3 = time.perf_counter()
        logging.info("Done {}, time: {}".format(each_file, t3 - t2))
    t4 = time.perf_counter()
    logging.info("All cost: {}".format(t4 - t1))


# nnunet_input -> nnunet_output
def nnunet_convert(from_path: str, to_path: str):
    os.makedirs(to_path, exist_ok=True)
    console_command = "nnUNet_predict -i {} -o {} -t 47 -m 2d -f 4".format(from_path, to_path)
    logging.info("Start to run command: {}".format(console_command))
    try:
        t1 = time.perf_counter()
        logging.info(os.system(console_command))
        t2 = time.perf_counter()
        logging.info("infer cost: {}".format(t2 - t1))
    except Exception as e:
        logging.error(traceback.format_exc())
        logging.error("An error occurred, please check the log")
        exit(-1)


# nnunet_output -> out_raw_data
def write_3d_to_2d(from_path: str, to_path: str, axis="z"):
    logging.info("Start to convert nii to image")

    os.makedirs(to_path, exist_ok=True)
    nii_file = [f for f in os.listdir(from_path) if f.lower().endswith(".nii.gz")]
    t1 = time.perf_counter()
    for each_nii in nii_file:
        t2 = time.perf_counter()
        name = each_nii[:-7]
        png_save_path = os.path.join(to_path, name + ".png")

        npy_data = np.asanyarray(nib.load(os.path.join(from_path, each_nii)).dataobj)
        npy_data = np.rot90(np.fliplr(npy_data))
        npy_data = npy_data[:, :, 0]
        cv2.imwrite(png_save_path, npy_data)

        t3 = time.perf_counter()
        logging.info("Done {}, time: {}".format(each_nii, t3 - t2))
    t4 = time.perf_counter()
    logging.info("All cost: {}".format(t4 - t1))


# out_raw_data -> out_color_data
def raw_to_color(from_path: str, to_path: str):
    logging.info("Start to draw color to raw output")

    os.makedirs(to_path, exist_ok=True)

    file_list = os.listdir(from_path)
    color_list = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (192, 255, 0), (0, 255, 156)]
    t1 = time.perf_counter()
    for each_file in file_list:
        t2 = time.perf_counter()
        name, _ = os.path.splitext(each_file)
        png_save_path = os.path.join(to_path, name + ".png")

        ori_img = cv2.imread(os.path.join(from_path, each_file))
        zero_img = np.zeros_like(ori_img, dtype=np.uint8)
        ori_img = ori_img[:, :, 0]

        for i in range(1, 6):
            zero_img[np.where(ori_img == i)] = color_list[i - 1]

        cv2.imwrite(png_save_path, zero_img)

        t3 = time.perf_counter()
        logging.info("Done {}, time: {}".format(each_file, t3 - t2))
    t4 = time.perf_counter()
    logging.info("All cost: {}".format(t4 - t1))


# out_color_data -> out_merge_data
def color_to_merge(ori_path: str, color_path: str, to_path: str):
    logging.info("Start to merge result")

    os.makedirs(to_path, exist_ok=True)

    file_list = os.listdir(ori_path)
    t1 = time.perf_counter()
    for each_file in file_list:
        t2 = time.perf_counter()
        name, _ = os.path.splitext(each_file)
        png_save_path = os.path.join(to_path, name + ".png")

        ori_img = cv2.imread(os.path.join(ori_path, each_file))
        color_img = cv2.imread(os.path.join(color_path, each_file))
        out_img = cv2.addWeighted(ori_img, 0.68, color_img, 0.32, 3)

        cv2.imwrite(png_save_path, out_img)

        t3 = time.perf_counter()
        logging.info("Done {}, time: {}".format(each_file, t3 - t2))
    t4 = time.perf_counter()
    logging.info("All cost: {}".format(t4 - t1))


# Function for performing vascular analysis
def perform_vascular_analysis(input_path: str, output_path: str):
    logging.info("Performing vascular analysis...")
    os.makedirs(output_path, exist_ok=True)

    file_list = os.listdir(input_path)
    for filename in file_list:
        if filename.endswith(".png"):
            input_image_path = os.path.join(input_path, filename)
            output_image_path = os.path.join(output_path, filename)

            # Read input image
            input_image = cv2.imread(input_image_path, cv2.IMREAD_GRAYSCALE)

            # Apply adaptive thresholding
            segmented_image = cv2.adaptiveThreshold(input_image, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                                    cv2.THRESH_BINARY_INV, 11, 2)

            # Save segmented image
            cv2.imwrite(output_image_path, segmented_image)

            logging.info(f"Vascular analysis completed for {filename}. Segmented image saved at {output_image_path}")

def run(base_path: str):
    t1 = time.perf_counter()
    # user_input -> nnunet_input
    write_2d_to_3d(os.path.join(base_path, "user_input"), os.path.join(base_path, "nnunet_input"))
    # nnunet_input -> nnunet_output
    nnunet_convert(os.path.join(base_path, "nnunet_input"), os.path.join(base_path, "nnunet_output"))
    # nnunet_output -> out_raw_data
    write_3d_to_2d(os.path.join(base_path, "nnunet_output"), os.path.join(base_path, "out_raw_data"))
    # out_raw_data -> out_color_data
    raw_to_color(os.path.join(base_path, "out_raw_data"), os.path.join(base_path, "out_color_data"))
    # Perform vascular analysis
    perform_vascular_analysis(os.path.join(base_path, "out_color_data"), os.path.join(base_path, "vascular_analysis_output"))
    # out_color_data -> out_merge_data
    color_to_merge(os.path.join(base_path, "user_input"), os.path.join(base_path, "out_color_data"),
                   os.path.join(base_path, "out_merge_data"))
    t2 = time.perf_counter()
    logging.info(
        "All cost {}, {} s/pic".format(t2 - t1, (t2 - t1) / len(os.listdir(os.path.join(base_path, "user_input"))) )
    )


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--workdir", type=str)
    args = parser.parse_args()
    if args.workdir == "":
        logging.error("work dir is empty")

    run(args.workdir)

